import hmac
import hashlib
import base64
import json
import time
import email.utils
import ssl
from urllib.parse import quote
from websocket import create_connection # type: ignore

class SparkAPI:
    def __init__(self, app_id, api_key, api_secret):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.host = "spark-api.xf-yun.com"  # 或根据区域改为 spark-api-gd.xf-yun.com
        self.path = "/v1.1/chat"  # 与控制台开通的版本一致

    def _generate_rfc1123_time(self):
        """生成RFC1123标准时间戳"""
        return email.utils.formatdate(time.time(), usegmt=True)

    def _generate_signature(self):
        """生成签名（修复打印问题）"""
        date = self._generate_rfc1123_time()
        signature_origin = (
            f"host: {self.host}\n"
            f"date: {date}\n"
            f"GET {self.path} HTTP/1.1"
        )
        print("[DEBUG] 签名原文:", repr(signature_origin))
        
        # 计算签名
        signature_sha = hmac.new(
            self.api_secret.encode("utf-8"),
            signature_origin.encode("utf-8"),
            digestmod=hashlib.sha256
        ).digest()
        signature = base64.b64encode(signature_sha).decode("utf-8")
        print("[DEBUG] 计算签名:", signature)  # 正确打印签名值
        return signature, date

    def _construct_auth_url(self):
        """构造鉴权URL"""
        signature, date = self._generate_signature()
        authorization_origin = (
            f'api_key="{self.api_key}", '
            f'algorithm="hmac-sha256", '
            f'headers="host date request-line", '
            f'signature="{signature}"'
        )
        authorization = base64.b64encode(authorization_origin.encode("utf-8")).decode("utf-8")
        return (
            f"wss://{self.host}{self.path}?"
            f"authorization={quote(authorization)}&"
            f"date={quote(date)}&"
            f"host={quote(self.host)}"
        )

    def get_response(self, prompt):
        try:
            ws_url = self._construct_auth_url()
            print("[DEBUG] 最终URL:", ws_url)
            
            # 建立带SSL验证的连接
            ws = create_connection(
                ws_url,
                sslopt={
                    "cert_reqs": ssl.CERT_REQUIRED,
                    "ca_certs": "/etc/ssl/certs/ca-certificates.crt"
                },
                timeout=30
            )
            
            # v2.1请求格式
            request_data = {
                "header": {"app_id": self.app_id},
                "parameter": {
                    "chat": {
                        "domain": "lite",
                        "temperature": 0.5,
                        "max_tokens": 4096
                    }
                },
                "payload": {
                    "message": {
                        "text": [{"role": "user", "content": prompt}]
                    }
                }
            }
            ws.send(json.dumps(request_data))
            
            # 处理响应
            full_response = ""
            while True:
                res = ws.recv()
                data = json.loads(res)
                if data["header"]["code"] != 0:
                    return f"Error {data['header']['code']}: {data['header']['message']}"
                full_response += data["payload"]["choices"]["text"][0]["content"]
                if data["header"]["status"] == 2:
                    break
            ws.close()
            return full_response
            
        except Exception as e:
            return f"系统错误: {str(e)}"